using System.Collections.Concurrent;

namespace Pinwheel.Kernel
{
    public class Rpc : IComponent, ISendHandler, IReceiveHandler
    {
        private ConcurrentDictionary<int, TaskCompletionSource<IMessage>> m_RpcMap =
            new ConcurrentDictionary<int, TaskCompletionSource<IMessage>>();
        private int m_Session = 0;

        public void OnAttach(IService service)
        {
            if (service is IMessageHandler mh)
            {
                mh.RegisterMessageHandler(MessageHandlerType.Send, HandleSend);
                mh.RegisterMessageHandler(MessageHandlerType.Receive, HandleReceive);
            }
        }

        public void OnDetach(IService service)
        {
            if (service is IMessageHandler mh)
            {
                mh.UnregisterMessageHandler(MessageHandlerType.Send, HandleSend);
                mh.UnregisterMessageHandler(MessageHandlerType.Receive, HandleReceive);
            }
        }

        public TaskCompletionSource<IMessage> HandleSend(ref ServiceMessage message)
        {
            var tcs = new TaskCompletionSource<IMessage>();
            var session = Interlocked.Increment(ref m_Session);
            if (session == 0) session = Interlocked.Increment(ref m_Session);
            message.SetSession(session);
            m_RpcMap.TryAdd(message.Session, tcs);
            return tcs;
        }

        public TaskCompletionSource<IMessage> HandleReceive(ref ServiceMessage message)
        {
            if (m_RpcMap.TryGetValue(message.Session, out var tcs))
            {
                tcs.SetResult(message.Message);
                m_RpcMap.TryRemove(message.Session, out _);
            }
            return tcs;
        }
    }
}